# -*- coding:utf-8 -*-
# import the necessary packages
#from keras.applications import ResNet50
from keras.applications import imagenet_utils
import numpy as np
import settings
import helpers
import redis
import time
import json

########### get current time #########
import datetime
def get_time_stamp():
    ct = time.time()
    local_time = time.localtime(ct)
    data_head = time.strftime("%Y-%m-%d %H:%M:%S", local_time)
    data_secs = (ct -int(ct)) * 1000
    time_stamp = "%s.%03d" % (data_head, data_secs)
    return time_stamp
##################################################
# connect to Redis server 连接到redis服务器
db = redis.StrictRedis(host=settings.REDIS_HOST,
    port=settings.REDIS_PORT, db=settings.REDIS_DB)

# import Mask RCNN
from mrcnn.config import Config
from mrcnn import model as modellib

######################
#  Configurations
######################
class nsfwConfig(Config):
    NAME = "nsfw"
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 1 + 6
    STEPS_PER_EPOCH = 100
    DETECTION_MIN_CONFIDENCE = 0.8

class InferenceConfig(nsfwConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
#####################
import os
import skimage.draw
from skimage import data,color
import skimage.transform
## import projective transform to change floor sketch
# import image_transform
###################################
project_path = os.path.dirname(os.path.realpath(__file__))
model_dirname = "models"
weights_filename = "mask_rcnn_nsfw_0494.h5"
weights_path = os.path.join(project_path , model_dirname , weights_filename )
model_path = weights_path

artimundi_path = os.path.join(project_path,"static", "products", "artimundi")
###################################
class_names = ['BG', 'butt', 'private', 'breast', 'animal', 'cartoon_kiss', 'cartoon_breast']


def classify_process():
    # load the pre-trained Keras model (here we are using a model
    # pre-trained on ImageNet and provided by Keras, but you can
    # substitute in your own networks just as easily) 加载模型，可以用自己的模型
    print("\n* Loading model...")
    config = InferenceConfig()
    config.display()
    ###### Create Model ######
    model = modellib.MaskRCNN(mode="inference", config=config, model_dir=model_path)
    ###### Load Weights ######
    model.load_weights(weights_path, by_name=True)
    print("\n** Model & Weights loaded")

    # continually pool for new images to classify 循环识别新图像
    while True:
        # attempt to grab a batch of images from the database, then
        # initialize the image IDs and batch of images themselves
        # 从数据库获取一组图片，初始化图片ID
        queue = db.lrange(settings.IMAGE_QUEUE, 0,
            settings.BATCH_SIZE - 1)
        imageIDs = []
        batch = None
 
        # loop over the queue 对图像数据库循环处理
        for q in queue:
            # deserialize the object and obtain the input image
            # 反序列化,获取输入图片
            q = json.loads(q.decode("utf-8"))
            img_in_width = int(q["width"])
            img_in_height = int(q["height"])
            # product_selected = q["product_selected"]
            # product_selected = image_transform.readname_txt(q["product_selected"]) + ".jpg"
            image = helpers.base64_decode_image(q["image"],
                settings.IMAGE_DTYPE,
                (img_in_height, img_in_width,
                 settings.IMAGE_CHANS))

            # check to see if the batch list is None
            if batch is None:
                batch = image
 
            # otherwise, stack the data
            else:
                batch = np.vstack([batch, image])
 
            # update the list of image IDs 上传图片ID列表
            imageIDs.append(q["id"])
 
        # check to see if we need to process the batch
        if len(imageIDs) > 0:
            # classify the batch
            print("** Batch size: {}\n".format(batch.shape))
            batch = batch.astype(np.uint8)
            # read product image from server

            # sketch = skimage.io.imread(product_selected)
            # sketch = image_transform.transform_sketch(sketch, 20)
            ###################
            r = model.detect([batch], verbose=1)[0]
            # splash = color_splash(batch, r['masks'])
            # splash = image_transform.generate_mask(batch, r['masks'], sketch)
            # print("\n***** splash finished !")
            # splashed_file_name = "splashed_{:%Y%m%dT%H%M%S}.png".format(datetime.datetime.now())
            # splashed_file_path = os.path.join(project_path, "static", "images", "results", splashed_file_name)
            # skimage.io.imsave(splashed_file_path, splash)
            # print("***** splash saved !\n")
            ###################
            # splashed_file_path = [splashed_file_name]
            results_class_ids = r["class_ids"]
            results_scores = r["scores"]
            print("\t ***** Class IDs: ", results_class_ids)
            print("\t ***** Scores: ", results_scores)
            # loop over the image IDs and their corresponding set of results from our model
            # for (imageID, class_id, score) in zip(imageIDs, results_class_ids, results_scores):
            for imageID in imageIDs:
                output = []
                if len(results_class_ids)<=0:

                    r = {"class": "no_result", "score": float(0)}
                    print("\t **** dict-r: ", r)
                    output.append(r)
                else:

                    for (class_id, score) in zip(results_class_ids, results_scores):

                        # initialize the list of output predictions

                        r = {"class": class_names[class_id], "score": score}
                        print("\t **** dict-r: ",r)
                        output.append(r)
                        # store the output predictions in the database, using
                        # the image ID as the key so we can fetch the results

                print("\t output: ", output)
                db.set(imageID, json.dumps(output))
                    
 
            # remove the set of images from our queue
            db.ltrim(settings.IMAGE_QUEUE, len(imageIDs), -1)
 
        # sleep for a small amount
        time.sleep(settings.SERVER_SLEEP)
 
# if this is the main thread of execution start the model server
# process 模型服务器主线程
if __name__ == "__main__":
    classify_process()
